package com.jessechanrui.template.common.template.auth.websocket;


import lombok.extern.slf4j.Slf4j;
import org.springframework.stereotype.Component;
import org.springframework.web.reactive.socket.WebSocketHandler;
import org.springframework.web.reactive.socket.WebSocketMessage;
import org.springframework.web.reactive.socket.WebSocketSession;
import reactor.core.publisher.Flux;
import reactor.core.publisher.Mono;
import reactor.core.publisher.Sinks;

import java.time.Duration;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;

@Component
@Slf4j
public class StreamReactiveWebSocketHandler implements WebSocketHandler {

    private final WebSocketSessionManager sessionManager;
    // 用于为每个会话创建独立的数据流
    private final Map<String, Sinks.Many<String>> sessionStreams = new ConcurrentHashMap<>();

    public StreamReactiveWebSocketHandler(WebSocketSessionManager sessionManager) {
        this.sessionManager = sessionManager;
    }

    @Override
    public Mono<Void> handle(WebSocketSession session) {
        String sessionId = extractSessionId(session.getHandshakeInfo().getUri().getPath());
        log.info("Reactive stream WebSocket connection established for session: {}", sessionId);

        // 注册会话
        sessionManager.registerReactiveSession(sessionId, session);

        // 为每个会话创建一个可多播的数据流
        Sinks.Many<String> sink = Sinks.many().multicast().onBackpressureBuffer();
        sessionStreams.put(sessionId, sink);

        // 欢迎消息
        Mono<WebSocketMessage> welcomeMessage = Mono.just(
                session.textMessage("Connected to reactive stream WebSocket. Session ID: " + sessionId +
                        "\nType 'stream:start' to begin receiving data stream" +
                        "\nType 'stream:stop' to stop the stream" +
                        "\nType 'stream:custom:N' to receive N custom messages")
        );

        // 处理客户端请求
        Flux<WebSocketMessage> response = session.receive()
                .map(WebSocketMessage::getPayloadAsText)
                .doOnNext(message -> handleClientCommand(sessionId, message))
                .map(message -> session.textMessage("Command received: " + message));

        // 合并欢迎消息、命令响应和数据流
        Flux<WebSocketMessage> output = welcomeMessage
                .concatWith(response)
                .concatWith(sink.asFlux().map(session::textMessage));

        // 连接关闭时清理资源
        return session.send(output)
                .doFinally(signalType -> {
                    sessionManager.removeReactiveSession(sessionId);
                    sessionStreams.remove(sessionId);
                    log.info("Reactive stream WebSocket connection closed for session: {}", sessionId);
                });
    }

    /**
     * 处理客户端命令
     */
    private void handleClientCommand(String sessionId, String command) {
        Sinks.Many<String> sink = sessionStreams.get(sessionId);
        if (sink == null) {
            log.warn("No stream found for session: {}", sessionId);
            return;
        }

        // 处理开始流命令
        if ("stream:start".equals(command)) {
            log.info("Starting continuous stream for session: {}", sessionId);

            // 创建一个持续的数据流（每秒发送一条消息）
            Flux.interval(Duration.ofSeconds(1))
                    .map(i -> "Stream data " + i + " at " + System.currentTimeMillis())
                    .subscribe(
                            data -> sink.tryEmitNext(data),
                            error -> log.error("Error in stream: {}", error.getMessage())
                    );
        }
        // 处理停止流命令
        else if ("stream:stop".equals(command)) {
            log.info("Stopping stream for session: {}", sessionId);
            // 关闭当前sink并创建新的sink，实现流的停止
            sessionStreams.remove(sessionId);
            sessionStreams.put(sessionId, Sinks.many().multicast().onBackpressureBuffer());
        }
        // 处理自定义数量的流数据
        else if (command.startsWith("stream:custom:")) {
            try {
                int count = Integer.parseInt(command.split(":")[2]);
                log.info("Sending {} custom stream messages to session: {}", count, sessionId);

                // 创建指定数量的数据流（每500ms一条）
                Flux.range(1, count)
                        .delayElements(Duration.ofMillis(500))
                        .map(i -> "Custom stream data " + i + " of " + count)
                        .subscribe(
                                data -> sink.tryEmitNext(data),
                                error -> log.error("Error in custom stream: {}", error.getMessage())
                        );
            } catch (Exception e) {
                sink.tryEmitNext("Invalid custom stream command. Use format: stream:custom:N");
            }
        }
        // 未知命令
        else if (command.startsWith("stream:")) {
            sink.tryEmitNext("Unknown command. Available commands: stream:start, stream:stop, stream:custom:N");
        }
    }

    /**
     * 从路径中提取会话ID
     */
    private String extractSessionId(String path) {
        String[] parts = path.split("/");
        return parts.length > 0 ? parts[parts.length - 1] : null;
    }

    /**
     * 外部触发向指定会话发送流数据
     */
    public void sendExternalStream(String sessionId, Flux<String> dataStream) {
        Sinks.Many<String> sink = sessionStreams.get(sessionId);
        if (sink != null) {
            dataStream.subscribe(
                    data -> sink.tryEmitNext(data),
                    error -> log.error("Error in external stream: {}", error.getMessage())
            );
        } else {
            log.warn("No active stream found for session: {}", sessionId);
        }
    }
}

